"""
Copyright (c) 2024 Beijing Jiaotong University
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Yunzhou Tang
Created: 2024/03/08
Supported by: National Key Research and Development Program of China
"""

import phot

if __name__ == "__main__":
    """ 双偏振光收发模块 + G.652光纤 """
    phot.config(plot=True, backend="torch")  # 配置基本信息 开启画图 使用torch作为运算库后端

    """ 设置全局系统仿真参数 """
    num_symbols = 2**16                              # 符号数目
    bits_per_symbol = 4                              # 2 for QPSK; 4 for 16QAM; 6 for 64QAM  设置调制格式
    total_baud = 40e9                                # 信号波特率，符号率
    up_sampling_factor = 4                           # 上采样倍数
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    Reference_frequency = 193.1e12

    """ 发射端 """
    """ RRC """
    RRC_ROLL_OFF = 0.01

    """ OSNR """
    osnr = 32

    """ power """
    power = 0     #入纤功率
    signals = phot.transmitter_B2B(
        num_symbols, bits_per_symbol, total_baud, up_sampling_factor, RRC_ROLL_OFF, osnr, power)

    """ 光纤信道"""
    num_spans = 1        # 多少个 span (每个span经过一次放大器)
    span_length = 50    # 一个 span 的长度 (km)
    delta_z = 1        # 单步步长 (km)
    alpha = 0.2
    beta2 = -21.6676e-24
    gamma = 1.3
    
    EDFA_nf = 5
    f_c = 193548387096774.16
    gain = 10.0

    for i in range(num_spans):
        signals = phot.fiber_G652(signals, num_symbols, bits_per_symbol, total_baud, up_sampling_factor, sampling_rate, num_spans, span_length ,delta_z, alpha, beta2, gamma)

        signals = phot.amplifier(signals, EDFA_nf, sampling_rate, f_c, gain) 

    """ 接收端 """
    """ CD """
    dispersion = 17
    Lambda = 1550
    FiberLength = num_spans * span_length
    beta2 = -21.6676e-24

    """ RRC """
    RRC_ROLL_OFF = 0.01

    """ 自适应均衡器 """    
    num_tap = 25
    ref_power_cma = 2
    cma_convergence = 33000
    step_size_cma = 2e-9
    step_size_rde = 2e-9

    """ BPS """
    num_test_angle = 64
    block_size = 120

    signals = phot.receiver_DSP(signals, num_symbols, bits_per_symbol, total_baud, up_sampling_factor,dispersion, Lambda, FiberLength, beta2, RRC_ROLL_OFF, num_tap, ref_power_cma, cma_convergence, step_size_cma, step_size_rde, num_test_angle, block_size)

    # 分析器画星座图
    phot.constellation_diagram(signals,isdata = False)

